#include <algorithm/math.h>
#include <algorithm/qekf.h>
#include <arm_math.h>

constexpr float ConstQuaternionEkfP[36] = {1e5, 0.1, 0.1, 0.1, 0.1, 0.1, //
                                           0.1, 1e5, 0.1, 0.1, 0.1, 0.1, //
                                           0.1, 0.1, 1e5, 0.1, 0.1, 0.1, //
                                           0.1, 0.1, 0.1, 1e5, 0.1, 0.1, //
                                           0.1, 0.1, 0.1, 0.1, 1e2, 0.1, //
                                           0.1, 0.1, 0.1, 0.1, 0.1, 1e2};

namespace os {
QuaternionExtendedKalmanFilter::QuaternionExtendedKalmanFilter(float process_noise1, float process_noise2,
                                                               float measure_noise, float lambda, float lpf)
    : P{ConstQuaternionEkfP},
      _acc_low_pass_filter_coef{lpf},
      _q1{process_noise1},
      _q2{process_noise2},
      _r{measure_noise},
      _chi_threshold{1e-8},
      _lambda{MIN(lambda, 1.0f)} {
    xHat = {1, 0, 0, 0, 0, 0};
}

QuaternionExtendedKalmanFilter* QuaternionExtendedKalmanFilter::update(float gx, float gy, float gz, float ax, float ay,
                                                                       float az, float dt) {
    static float half_gx_dt;
    static float half_gy_dt;
    static float half_gz_dt;
    static float accel_inv_norm;

    _gyro[0] = gx - xHat[4][0];
    _gyro[1] = gy - xHat[5][0];
    _gyro[2] = gz - 0;

    half_gx_dt = 0.5F * _gyro[0] * dt;
    half_gy_dt = 0.5F * _gyro[1] * dt;
    half_gz_dt = 0.5F * _gyro[2] * dt;

    /*   状态转移矩阵 F, # 为需要修改的值
     *        [0]    [1]    [2]    [3]    [4]    [5]
     * [0]     1      #      #      #      0      0
     * [1]     #      1      #      #      0      0
     * [2]     #      #      1      #      0      0
     * [3]     #      #      #      1      0      0
     * [4]     0      0      0      0      1      0
     * [5]     0      0      0      0      0      1
     */
    F = MatrixI<6>;

    F[0][1] = -half_gx_dt;
    F[0][2] = -half_gy_dt;
    F[0][3] = -half_gz_dt;

    F[1][0] = +half_gx_dt;
    F[1][2] = +half_gz_dt;
    F[1][3] = -half_gy_dt;

    F[2][0] = +half_gy_dt;
    F[2][1] = -half_gz_dt;
    F[2][3] = +half_gx_dt;

    F[3][0] = +half_gz_dt;
    F[3][1] = +half_gy_dt;
    F[3][2] = -half_gx_dt;

    if (!_update_count) {
        _accel[0] = ax;
        _accel[1] = ay;
        _accel[2] = az;
    }
    _accel[0] = (_accel[0] * _acc_low_pass_filter_coef + ax * dt) / (dt + _acc_low_pass_filter_coef);
    _accel[1] = (_accel[1] * _acc_low_pass_filter_coef + ay * dt) / (dt + _acc_low_pass_filter_coef);
    _accel[2] = (_accel[2] * _acc_low_pass_filter_coef + az * dt) / (dt + _acc_low_pass_filter_coef);

    accel_inv_norm = os::math::invSqrt(_accel[0] * _accel[0] + _accel[1] * _accel[1] + _accel[2] * _accel[2]);
    for (uint8_t i = 0; i < decltype(z)::ROW; i++) {
        z[i][0] = _accel[i] * accel_inv_norm;
    }

    _gyro_norm = 1.0F / os::math::invSqrt(_gyro[0] * _gyro[0] + _gyro[1] * _gyro[1] + _gyro[2] * _gyro[2]);
    _accl_norm = 1.0F / accel_inv_norm;

    if (_gyro_norm < 0.3F && _accl_norm > 9.8F - 0.5F && _accl_norm < 9.8F + 0.5F) {
        _stable_flag = 1;
    } else {
        _stable_flag = 0;
    }

    /*   噪声矩阵 Q R, # 为需要修改的值
     *  a = _q1 * dt, b = _q2 * dt
     *        [0]    [1]    [2]    [3]    [4]    [5]
     * [0]     a      0      0      0      0      0
     * [1]     0      a      0      0      0      0
     * [2]     0      0      a      0      0      0
     * [3]     0      0      0      a      0      0
     * [4]     0      0      0      0      b      0
     * [5]     0      0      0      0      0      b
     * R: r = _r
     *        [0]    [1]    [2]
     * [0]     r      0      0
     * [1]     0      r      0
     * [2]     0      0      r
     */

    Q[0][0] = _q1 * dt;
    Q[1][1] = _q1 * dt;
    Q[2][2] = _q1 * dt;
    Q[3][3] = _q1 * dt;
    Q[4][4] = _q2 * dt;
    Q[5][5] = _q2 * dt;

    R[0][0] = _r;
    R[0][4] = _r;
    R[0][8] = _r;

    // 1. xhat'(k)= A·xhat(k-1) + B·u
    xHat_ = F * xHat;

    static float* q = xHat_[0];
    static float qInvNorm;

    // 四元数归一化
    qInvNorm = os::math::invSqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]);
    for (uint8_t i = 0; i < 4; i++) {
        xHat_[i][0] *= qInvNorm;
    }
    /*   状态转移矩阵 F, # 为需要修改的值
     *        [0]    [1]    [2]    [3]    [4]    [5]
     * [0]     1      -      -      -      #      #
     * [1]     -      1      -      -      #      #
     * [2]     -      -      1      -      #      #
     * [3]     -      -      -      1      #      #
     * [4]     0      0      0      0      1      0
     * [5]     0      0      0      0      0      1
     */
    F[0][4] = +q[1] * dt / 2;
    F[0][5] = +q[2] * dt / 2;
    F[1][4] = -q[0] * dt / 2;
    F[1][5] = +q[3] * dt / 2;
    F[2][4] = -q[3] * dt / 2;
    F[2][5] = -q[0] * dt / 2;
    F[3][4] = +q[2] * dt / 2;
    F[3][5] = -q[1] * dt / 2;

    // 渐消
    P[4][4] = MIN(P[4][4] / _lambda, 10000);
    P[5][5] = MIN(P[5][5] / _lambda, 10000);

    // 2. P'(k) = A·P(k-1)·AT + Q
    P_ = F * P * F.T() + Q;

    /* H:
     *        [0]    [1]    [2]    [3]    [4]    [5]
     * [0]   -2q2    2q3   -2q0    2q1     -      -
     * [1]    2q1    2q0    2q3    2q2     -      -
     * [2]    2q0   -2q1   -2q2    2q3     -      -
     */
    H[0][0] = -q[2];
    H[0][1] = +q[3];
    H[0][2] = -q[0];
    H[0][3] = +q[1];

    H[1][0] = +q[1];
    H[1][1] = +q[0];
    H[1][2] = +q[3];
    H[1][3] = +q[2];

    H[2][0] = +q[0];
    H[2][1] = -q[1];
    H[2][2] = -q[2];
    H[2][3] = +q[3];

    H = H * 2;

    // 3. K(k) = P'(k)·HT / (H·P'(k)·HT + R)
    static Matrix<3, 3> temp_matrix;
    temp_matrix = (H * P_ * (H.T()) + R).inverse();

    static Matrix<3, 1> temp_vector;
    temp_vector = {
        2 * (q[1] * q[3] - q[0] * q[2]), //
        2 * (q[0] * q[1] + q[2] * q[3]), //
        q[0] * q[0] - q[1] * q[1] - q[2] * q[2] + q[3] * q[3] //
    };

    for (uint8_t i = 0; i < 3; i++) {
        _orientation_cosine[i] = acosf(fabsf(temp_vector[0][i]));
    }
    // 卡方校验
    static Matrix<3, 1> temp_vector1;
    temp_vector1 = z - temp_vector;
    _chi = (temp_vector1.T()) * (temp_matrix * temp_vector1);

    // Rk很小，滤波收敛/收敛
    if (_chi[0][0] < 0.5F * _chi_threshold) {
        _converge_flag = true;
    }
    // Rk大于3，但上一次是收敛的
    if (_chi[0][0] > _chi_threshold && _converge_flag) {
        _error_count += _stable_flag;
        if (_error_count <= 50) {
            xHat = xHat_;
            _normal_p_update = false;
            return this;
        }
        _normal_p_update = true;
        _converge_flag = false;
    } else {
        // 如果发散或 rk 不是那么大/可接受，使用自适应增益
        if (_chi[0][0] > 0.1F * _chi_threshold && _converge_flag) {
            _adaptive_gain_scale = (_chi_threshold - _chi[0][0]) / (0.9f * _chi_threshold);
        } else {
            _adaptive_gain_scale = 1;
        }
        _error_count = 0;
        _normal_p_update = true;
    }

    // cal kf-gain K
    K = P_ * (H.T()) * temp_matrix;

    // implement adaptive
    for (uint32_t i = 0; i < decltype(K)::ROW * decltype(K)::COL; i++) {
        K[0][i] *= _adaptive_gain_scale;
    }
    for (uint32_t i = 4; i < decltype(K)::ROW; i++) {
        for (uint8_t j = 0; j < decltype(K)::COL; j++) {
            K[i][j] *= _orientation_cosine[i - 4] / 1.5707963F; // 1 rad
        }
    }

    static Matrix<6, 1> tempVector2;
    tempVector2 = K * (z - temp_vector);

    if (_converge_flag) {
        for (uint8_t i = 4; i < decltype(tempVector2)::ROW; i++) {
            if (tempVector2[i][0] > 1e-2F * dt) {
                tempVector2[i][0] = 1e-2F * dt;
            }
            if (tempVector2[i][0] < -1e-2F * dt) {
                tempVector2[i][0] = -1e-2F * dt;
            }
        }
    }

    tempVector2[3][0] = 0;

    xHat = xHat_ + tempVector2;

    // 5. P(k) = (1-K(k)·H)·P'(k) ==> P(k) = P'(k)-K(k)·H·P'(k)
    if (_normal_p_update) {
        P = P_ - (K * H * P_);
    } else {
        P = P_;
    }

    _update_count++;
    return this;
}
} // namespace os

namespace os {
namespace math {
Quaternion EulerToQuaternions(EulerAngles euler) {
    // Abbreviations for the various angular functions
    float32_t cy = cos(euler.yaw * 0.5);
    float32_t sy = sin(euler.yaw * 0.5);
    float32_t cp = cos(euler.pitch * 0.5);
    float32_t sp = sin(euler.pitch * 0.5);
    float32_t cr = cos(euler.roll * 0.5);
    float32_t sr = sin(euler.roll * 0.5);

    Quaternion q;
    q.w = cy * cp * cr + sy * sp * sr;
    q.x = cy * cp * sr - sy * sp * cr;
    q.y = sy * cp * sr + cy * sp * cr;
    q.z = sy * cp * cr - cy * sp * sr;

    return q;
}

EulerAngles QuaternionsToEuler(Quaternion q) {
    EulerAngles euler;

    // roll (x-axis rotation)
    euler.roll = atan2(2 * (q.w * q.x + q.y * q.z), 1 - 2 * (q.x * q.x + q.y * q.y));

    // pitch (y-axis rotation)
    float32_t sinp = 2 * (q.w * q.y - q.z * q.x);
    if (fabs(sinp) >= 1) {
        euler.pitch = radian_t(PI / 2 * sign(sinp)); // use 90 degrees if out of range
    } else {
        euler.pitch = asin(sinp);
    }

    // yaw (z-axis rotation)
    euler.yaw = atan2(2 * (q.w * q.z + q.x * q.y), 1 - 2 * (q.y * q.y + q.z * q.z));

    return euler;
}

void CorrectEulerAngles(EulerAngles* res, EulerAngles const& last) {
    using namespace units::literals;
    radian_t diff = last.roll - res->roll;
    radian_t delta = 0_deg;
    if (diff > 270_deg) {
        delta = 360_deg;
    } else if (diff > 90_deg) {
        delta = 180_deg;
    } else if (diff > -90_deg) {
        delta = 0_deg;
    } else if (diff > -270_deg) {
        delta = -180_deg;
    } else {
        delta = -360_deg;
    }
    res->roll += delta;

    diff = last.yaw - res->yaw;
    if (diff > 270_deg) {
        delta = 360_deg;
    } else if (diff > 90_deg) {
        delta = 180_deg;
    } else if (diff > -90_deg) {
        delta = 0_deg;
    } else if (diff > -270_deg) {
        delta = -180_deg;
    } else {
        delta = -360_deg;
    }
    res->yaw += delta;
}
} // namespace math
} // namespace os